import os
import sys
import time
import requests
from typing import Dict, Any, List
from dotenv import load_dotenv
from datetime import datetime

load_dotenv()


class Prometheus:
    def __init__(self, base_url: str, timeout: int = 10):
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout
        self.session = requests.Session()

    def query(self, expr: str) -> List[Dict[str, Any]]:
        url = f"{self.base_url}/api/v1/query"
        resp = self.session.get(url, params={"query": expr}, timeout=self.timeout)
        resp.raise_for_status()
        payload = resp.json()
        if payload.get("status") != "success":
            raise RuntimeError(f"Query failed: {payload}")
        return payload["data"]["result"]


def extract_instances(up_results: List[Dict[str, Any]]) -> List[str]:
    instances = []
    for item in up_results:
        inst = item.get("metric", {}).get("instance")
        if inst:
            instances.append(inst)
    return sorted(set(instances))


def value_to_float(val) -> float:
    try:
        # Instant vector value is like ["<timestamp>", "<value>"] or just "<value>"
        if isinstance(val, list):
            return float(val[1])
        return float(val)
    except Exception:
        return float("nan")


def generate_html_report(now: str, prom_url: str, job: str, data: List[Dict[str, Any]]) -> str:
    """生成HTML报表"""
    html = f"""
    <!DOCTYPE html>
    <html>
    <head>
        <meta charset="UTF-8">
        <title>节点巡检报表 - {now}</title>
        <style>
            body {{ font-family: Arial, sans-serif; margin: 20px; }}
            table {{ border-collapse: collapse; width: 100%; }}
            th, td {{ border: 1px solid #ddd; padding: 8px; text-align: left; }}
            th {{ background-color: #f2f2f2; }}
            tr:nth-child(even) {{ background-color: #f9f9f9; }}
            .header {{ margin-bottom: 20px; }}
            .warning {{ background-color: #ffeb3b; }}
            .critical {{ background-color: #ff9999; }}
        </style>
    </head>
    <body>
        <div class="header">
            <h2>节点巡检报表</h2>
            <p>巡检时间: {now}</p>
            <p>Prometheus: {prom_url} | Job: {job}</p>
        </div>
        <table>
            <tr>
                <th>Instance(IP:port)</th>
                <th>IP</th>
                <th>内存使用率(%)</th>
                <th>磁盘使用率(%)</th>
                <th>CPU使用率(%)</th>
                <th>网络接收(bps)</th>
                <th>网络发送(bps)</th>
                <th>TCP连接数</th>
                <th>Load1</th>
                <th>Load5</th>
                <th>Load15</th>
            </tr>
    """

    for item in data:
        # 添加警告颜色（示例：内存或磁盘使用率>80%为warning，>90%为critical）
        mem_class = "critical" if item['mem'] > 90 else "warning" if item['mem'] > 80 else ""
        disk_class = "critical" if item['disk'] > 90 else "warning" if item['disk'] > 80 else ""
        cpu_class = "critical" if item['cpu'] > 90 else "warning" if item['cpu'] > 80 else ""

        html += f"""
            <tr>
                <td>{item['instance']}</td>
                <td>{item['ip']}</td>
                <td class="{mem_class}">{item['mem']:.2f}</td>
                <td class="{disk_class}">{item['disk']:.2f}</td>
                <td class="{cpu_class}">{item['cpu']:.2f}</td>
                <td>{item['network_receive']:.2f}</td>
                <td>{item['network_transmit']:.2f}</td>
                <td>{item['tcp_conn']:.0f}</td>
                <td>{item['load1']:.2f}</td>
                <td>{item['load5']:.2f}</td>
                <td>{item['load15']:.2f}</td>
            </tr>
        """

    html += """
        </table>
    </body>
    </html>
    """
    return html


def main():
    prom_url = os.getenv("PROM_URL", "https://prometheus.local.com/")
    job = os.getenv("PROM_JOB", "node-nok8s")
    prom = Prometheus(prom_url)

    # 1) 获取实例列表（基于 up 向量）-
    up_expr = f'up{{job="{job}"}}'
    up_vec = prom.query(up_expr)
    instances = extract_instances(up_vec)
    if not instances:
        print(f"未发现实例（用表达式 {up_expr}）。请检查 Prometheus 抓取或 job 标签。")
        sys.exit(1)

    # 2) 组装多条 PromQL，一次性取回结果并映射
    mem_expr = (
        f'(1 - (node_memory_MemAvailable_bytes{{job="{job}"}} '
        f'/ node_memory_MemTotal_bytes{{job="{job}"}})) * 100'
    )

    # 过滤掉临时/容器/只读文件系统，按实例汇总
    disk_filter = 'fstype!~"tmpfs|overlay|squashfs|aufs"'
    disk_expr = (
        f'(1 - (sum by (instance) (node_filesystem_avail_bytes{{job="{job}",{disk_filter}}}) '
        f'/ sum by (instance) (node_filesystem_size_bytes{{job="{job}",{disk_filter}}}))) * 100'
    )

    load1_expr = f'node_load1{{job="{job}"}}'
    load5_expr = f'node_load5{{job="{job}"}}'
    load15_expr = f'node_load15{{job="{job}"}}'

    # 添加新的 PromQL 表达式
    cpu_expr = f'100 - (avg by (instance) (irate(node_cpu_seconds_total{{job="{job}",mode="idle"}}[5m])) * 100)'
    network_receive_expr = f'irate(node_network_receive_bytes_total{{job="{job}",device!~"lo|veth.*|docker.*|br.*|flannel.*"}}[5m])*8'
    network_transmit_expr = f'irate(node_network_transmit_bytes_total{{job="{job}",device!~"lo|veth.*|docker.*|br.*|flannel.*"}}[5m])*8'
    tcp_conn_expr = f'node_netstat_Tcp_CurrEstab{{job="{job}"}}'

    results = {
        "mem_used_pct": prom.query(mem_expr),
        "disk_used_pct": prom.query(disk_expr),
        "cpu_used_pct": prom.query(cpu_expr),
        "network_receive": prom.query(network_receive_expr),
        "network_transmit": prom.query(network_transmit_expr),
        "tcp_conn": prom.query(tcp_conn_expr),
        "load1": prom.query(load1_expr),
        "load5": prom.query(load5_expr),
        "load15": prom.query(load15_expr),
    }

    # 3) 将结果按 instance 聚合
    def map_by_instance(vec):
        m = {}
        for item in vec:
            inst = item.get("metric", {}).get("instance")
            if not inst:
                continue
            m[inst] = value_to_float(item.get("value"))
        return m

    mem_map = map_by_instance(results["mem_used_pct"])
    disk_map = map_by_instance(results["disk_used_pct"])
    cpu_map = map_by_instance(results["cpu_used_pct"])
    net_recv_map = map_by_instance(results["network_receive"])
    net_send_map = map_by_instance(results["network_transmit"])
    tcp_conn_map = map_by_instance(results["tcp_conn"])
    l1_map = map_by_instance(results["load1"])
    l5_map = map_by_instance(results["load5"])
    l15_map = map_by_instance(results["load15"])

    # 4) 打印
    now = time.strftime("%Y-%m-%d %H:%M:%S")
    print(f"时间: {now}")
    print(f"Prometheus: {prom_url}  | job: {job}")
    print("-" * 80)
    print(
        f"{'Instance(IP:port)':<28}  {'IP':<16}  {'Mem%':>7}  {'Disk%':>7}  {'CPU%':>7}  {'NetRecv':>10}  {'NetSend':>10}  {'TCP':>7}  {'Load1':>7}  {'Load5':>7}  {'Load15':>7}")
    print("-" * 140)

    report_data = []
    for inst in instances:
        ip = inst.split(":")[0]  # 若是主机名，这里会显示主机名；需要纯 IP 可在抓取端保证 instance=IP:port
        mem = mem_map.get(inst, float("nan"))
        dsk = disk_map.get(inst, float("nan"))
        cpu = cpu_map.get(inst, float("nan"))
        net_recv = net_recv_map.get(inst, float("nan"))
        net_send = net_send_map.get(inst, float("nan"))
        tcp_conn = tcp_conn_map.get(inst, float("nan"))
        l1 = l1_map.get(inst, float("nan"))
        l5 = l5_map.get(inst, float("nan"))
        l15 = l15_map.get(inst, float("nan"))

        def fmt(x):
            return "NaN" if x != x else f"{x:.2f}"

        # 控制台输出
        print(
            f"{inst:<28}  {ip:<16}  {fmt(mem):>7}  {fmt(dsk):>7}  {fmt(cpu):>7}  {fmt(net_recv):>10}  {fmt(net_send):>10}  {fmt(tcp_conn):>7}  {fmt(l1):>7}  {fmt(l5):>7}  {fmt(l15):>7}")

        # 收集报表数据
        report_data.append({
            'instance': inst,
            'ip': ip,
            'mem': mem if mem == mem else 0,
            'disk': dsk if dsk == dsk else 0,
            'cpu': cpu if cpu == cpu else 0,
            'network_receive': net_recv if net_recv == net_recv else 0,
            'network_transmit': net_send if net_send == net_send else 0,
            'tcp_conn': tcp_conn if tcp_conn == tcp_conn else 0,
            'load1': l1 if l1 == l1 else 0,
            'load5': l5 if l5 == l5 else 0,
            'load15': l15 if l15 == l15 else 0
        })

    # 生成HTML报表
    now = time.strftime("%Y-%m-%d_%H-%M-%S")
    report_html = generate_html_report(now, prom_url, job, report_data)

    # 保存HTML报表
    filename = f"linux_node_check_{now}.html"
    with open(filename, "w", encoding="utf-8") as f:
        f.write(report_html)

    print(f"\n报表已生成: {filename}")


if __name__ == "__main__":
    main()
